package com.mini.grpc.transport.netty;

import com.mini.grpc.common.Metadata;
import com.mini.grpc.common.Status;
import com.mini.grpc.serialization.Marshaller;
import com.mini.grpc.transport.ServerStream;
import com.mini.grpc.transport.ServerStreamListener;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.http2.DefaultHttp2DataFrame;
import io.netty.handler.codec.http2.DefaultHttp2Headers;
import io.netty.handler.codec.http2.DefaultHttp2HeadersFrame;
import io.netty.handler.codec.http2.Http2ConnectionEncoder;
import io.netty.handler.codec.http2.Http2Headers;
import io.netty.util.AsciiString;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;

/**
 * NettyServerStream实现了基于Netty的服务端流。
 *
 * @param <RespT> 响应类型
 */
public class NettyServerStream<RespT> implements ServerStream<RespT> {
    
    private static final Logger logger = LoggerFactory.getLogger(NettyServerStream.class);
    
    private final int streamId;
    private final NettyServerHandler handler;
    private final Http2ConnectionEncoder encoder;
    private ServerStreamListener<RespT> listener;
    private boolean headersSent = false;
    private Marshaller<RespT> marshaller; // 实际使用时需要从方法描述符获取
    
    /**
     * 创建一个新的NettyServerStream
     *
     * @param streamId 流ID
     * @param handler 服务端处理器
     * @param encoder HTTP/2连接编码器
     */
    public NettyServerStream(int streamId, NettyServerHandler handler, Http2ConnectionEncoder encoder) {
        this.streamId = streamId;
        this.handler = handler;
        this.encoder = encoder;
    }
    
    /**
     * 设置流监听器
     *
     * @param listener 流监听器
     */
    public void setListener(ServerStreamListener<RespT> listener) {
        this.listener = listener;
    }
    
    /**
     * 设置响应序列化器
     *
     * @param marshaller 响应序列化器
     */
    public void setMarshaller(Marshaller<RespT> marshaller) {
        this.marshaller = marshaller;
    }
    
    @Override
    public void sendHeaders(Metadata headers) {
        if (headersSent) {
            return;
        }
        
        Http2Headers http2Headers = new DefaultHttp2Headers();
        http2Headers.status(AsciiString.of("200"));
        http2Headers.set(AsciiString.of("content-type"), AsciiString.of("application/grpc"));
        
        // 添加自定义头部
        for (String key : headers.keys()) {
            if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
                // 二进制头部
                byte[] value = headers.get(Metadata.Key.ofBinary(key));
                if (value != null) {
                    http2Headers.set(AsciiString.of(key), AsciiString.of(new String(value, StandardCharsets.UTF_8)));
                }
            } else {
                // ASCII头部
                String value = headers.get(Metadata.Key.of(key));
                if (value != null) {
                    http2Headers.set(AsciiString.of(key), AsciiString.of(value));
                }
            }
        }
        
        encoder.writeHeaders(handler.getContext(), streamId, http2Headers, 0, false, handler.getContext().newPromise());
        handler.getContext().flush();
        headersSent = true;
    }
    
    @Override
    public void sendMessage(RespT message) {
        if (!headersSent) {
            sendHeaders(new Metadata());
        }
        
        if (marshaller == null) {
            logger.error("Marshaller not set for stream {}", streamId);
            return;
        }
        
        try {
            // 序列化消息
            InputStream messageStream = marshaller.stream(message);
            byte[] messageBytes = readAllBytes(messageStream);
            
            // gRPC消息格式：
            // 1字节压缩标志 + 4字节长度 + 消息内容
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            
            // 压缩标志（1字节，0表示不压缩）
            baos.write(0);
            
            // 消息长度（4字节，大端序）
            int length = messageBytes.length;
            baos.write((length >> 24) & 0xFF);
            baos.write((length >> 16) & 0xFF);
            baos.write((length >> 8) & 0xFF);
            baos.write(length & 0xFF);
            
            // 消息内容
            baos.write(messageBytes);
            
            // 发送数据
            ByteBuf data = Unpooled.wrappedBuffer(baos.toByteArray());
            logger.debug("Sending message for stream {}, size: {}", streamId, data.readableBytes());
            
            encoder.writeData(handler.getContext(), streamId, data, 0, false, handler.getContext().newPromise());
            handler.getContext().flush();
        } catch (IOException e) {
            logger.error("Failed to send message for stream {}", streamId, e);
        }
    }
    
    @Override
    public void close(Status status, Metadata trailers) {
        if (!headersSent) {
            // 如果头部未发送，将状态码放在头部
            Http2Headers headers = new DefaultHttp2Headers();
            headers.status(AsciiString.of(String.valueOf(grpcStatusToHttpStatus(status))));
            headers.set(AsciiString.of("grpc-status"), AsciiString.of(String.valueOf(status.getCode().value())));
            
            if (status.getDescription() != null) {
                headers.set(AsciiString.of("grpc-message"), AsciiString.of(status.getDescription()));
            }
            
            // 添加尾部元数据
            for (String key : trailers.keys()) {
                if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
                    // 二进制尾部
                    byte[] value = trailers.get(Metadata.Key.ofBinary(key));
                    if (value != null) {
                        headers.set(AsciiString.of(key), AsciiString.of(new String(value, StandardCharsets.UTF_8)));
                    }
                } else {
                    // ASCII尾部
                    String value = trailers.get(Metadata.Key.of(key));
                    if (value != null) {
                        headers.set(AsciiString.of(key), AsciiString.of(value));
                    }
                }
            }
            
            encoder.writeHeaders(handler.getContext(), streamId, headers, 0, true, handler.getContext().newPromise());
        } else {
            // 如果头部已发送，将状态码放在尾部
            Http2Headers trailerHeaders = new DefaultHttp2Headers();
            trailerHeaders.set(AsciiString.of("grpc-status"), AsciiString.of(String.valueOf(status.getCode().value())));
            
            if (status.getDescription() != null) {
                trailerHeaders.set(AsciiString.of("grpc-message"), AsciiString.of(status.getDescription()));
            }
            
            // 添加尾部元数据
            for (String key : trailers.keys()) {
                if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
                    // 二进制尾部
                    byte[] value = trailers.get(Metadata.Key.ofBinary(key));
                    if (value != null) {
                        trailerHeaders.set(AsciiString.of(key), AsciiString.of(new String(value, StandardCharsets.UTF_8)));
                    }
                } else {
                    // ASCII尾部
                    String value = trailers.get(Metadata.Key.of(key));
                    if (value != null) {
                        trailerHeaders.set(AsciiString.of(key), AsciiString.of(value));
                    }
                }
            }
            
            encoder.writeHeaders(handler.getContext(), streamId, trailerHeaders, 0, true, handler.getContext().newPromise());
        }
        
        handler.getContext().flush();
        handler.removeStream(streamId);
    }
    
    @Override
    public void request(int numMessages) {
        // 在服务端实现中，这个方法不需要做任何事情
    }
    
    @Override
    public int getStreamId() {
        return streamId;
    }
    
    @Override
    public void cancel() {
        // 取消流，发送RST_STREAM帧
        try {
            encoder.writeRstStream(handler.getContext(), streamId, 8, handler.getContext().newPromise());
            handler.getContext().flush();
        } catch (Exception e) {
            logger.error("Failed to cancel stream {}", streamId, e);
        } finally {
            handler.removeStream(streamId);
        }
    }
    
    /**
     * 处理接收到的数据
     *
     * @param data 数据
     * @param endOfStream 是否是流的最后一个数据帧
     */
    public void transportDataReceived(ByteBuf data, boolean endOfStream) {
        logger.debug("=== transportDataReceived called: streamId={}, dataSize={}, endOfStream={}, listener={} ===", 
                    streamId, data.readableBytes(), endOfStream, listener != null ? "SET" : "NULL");
        
        if (listener == null) {
            logger.error("Listener not set for stream {}", streamId);
            return;
        }
        
        try {
            logger.debug("Received data for stream {}, size: {}", streamId, data.readableBytes());
            
            // 解析gRPC消息
            // 格式：压缩标志(1字节) + 消息长度(4字节) + 消息内容
            
            while (data.isReadable()) {
                // 确保有足够的数据读取头部
                if (data.readableBytes() < 5) {
                    logger.debug("Not enough data for message header, waiting for more data");
                    return;
                }
                
                // 标记读取位置，以便回滚
                int readerIndex = data.readerIndex();
                
                // 读取压缩标志
                boolean compressed = data.readByte() != 0;
                
                // 读取消息长度（大端序）
                int length = data.readInt();
                
                logger.debug("Message header: compressed={}, length={}", compressed, length);
                
                // 确保有足够的数据读取消息体
                if (data.readableBytes() < length) {
                    // 不完整的消息，回滚读取位置并等待更多数据
                    data.readerIndex(readerIndex);
                    logger.debug("Not enough data for message body, waiting for more data");
                    return;
                }
                
                // 读取消息内容
                byte[] messageBytes = new byte[length];
                data.readBytes(messageBytes);
                
                logger.debug("Read message: {} bytes", messageBytes.length);
                
                // 反序列化消息
                // 这里我们使用StringMarshaller进行反序列化
                try {
                    ByteArrayInputStream bais = new ByteArrayInputStream(messageBytes);
                    @SuppressWarnings("unchecked")
                    RespT message = (RespT) new String(messageBytes, StandardCharsets.UTF_8);
                    
                    logger.debug("Deserialized message: {}", message);
                    
                    // 通知监听器
                    listener.messageReceived(message);
                } catch (Exception e) {
                    logger.error("Failed to deserialize message for stream {}", streamId, e);
                    transportReportStatus(Status.INTERNAL.withDescription("Failed to deserialize message: " + e.getMessage()));
                    return;
                }
            }
        } catch (Exception e) {
            logger.error("Error processing message for stream {}", streamId, e);
            transportReportStatus(Status.INTERNAL.withDescription("Error processing message: " + e.getMessage()));
        }
    }
    
    /**
     * 处理流半关闭
     */
    public void transportHalfClosed() {
        if (listener != null) {
            listener.halfClosed();
        }
    }
    
    /**
     * 处理收到的尾部头部
     *
     * @param trailers 尾部元数据
     */
    public void transportTrailersReceived(Metadata trailers) {
        // 在服务端实现中，这个方法不需要做任何事情
    }
    
    /**
     * 报告状态
     *
     * @param status 状态
     */
    public void transportReportStatus(Status status) {
        if (listener != null) {
            listener.cancelled(status);
        }
        handler.removeStream(streamId);
    }
    
    /**
     * 将gRPC状态码转换为HTTP状态码
     *
     * @param status gRPC状态
     * @return HTTP状态码
     */
    private int grpcStatusToHttpStatus(Status status) {
        switch (status.getCode()) {
            case OK:
                return 200;
            case CANCELLED:
                return 499;
            case UNKNOWN:
                return 500;
            case INVALID_ARGUMENT:
                return 400;
            case DEADLINE_EXCEEDED:
                return 504;
            case NOT_FOUND:
                return 404;
            case ALREADY_EXISTS:
                return 409;
            case PERMISSION_DENIED:
                return 403;
            case RESOURCE_EXHAUSTED:
                return 429;
            case FAILED_PRECONDITION:
                return 400;
            case ABORTED:
                return 409;
            case OUT_OF_RANGE:
                return 400;
            case UNIMPLEMENTED:
                return 501;
            case INTERNAL:
                return 500;
            case UNAVAILABLE:
                return 503;
            case DATA_LOSS:
                return 500;
            case UNAUTHENTICATED:
                return 401;
            default:
                return 500;
        }
    }
    
    /**
     * 读取输入流中的所有字节
     *
     * @param input 输入流
     * @return 字节数组
     * @throws IOException 如果发生I/O错误
     */
    private byte[] readAllBytes(InputStream input) throws IOException {
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        int n;
        byte[] data = new byte[1024];
        while ((n = input.read(data, 0, data.length)) != -1) {
            buffer.write(data, 0, n);
        }
        return buffer.toByteArray();
    }
} 